#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
作者: gavin
创建日期: 2024-10-17
更新日期: 2024-10-17
版本: 1.0

简要描述: 读取文件，挺替换文件内容

许可证:
这个文件是在CC BY-SA 4.0许可下发布的。
"""
import os
import re
import requests
from PIL import Image
import io

def download_image(url):
    """
    下载图片并返回图片数据。
    
    :param url: 图片URL
    :return: 图片数据或None（如果下载失败）
    """
    try:
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }
        response = requests.get(url, headers=headers, timeout=30)
        response.raise_for_status()
        print(f"成功下载图片: {url}")
        return response.content
    except Exception as e:
        print(f"下载图片失败 {url}: {str(e)}")
        return None

def get_file_extension(image_data):
    """
    获取图片的文件扩展名。
    
    :param image_data: 图片数据（字节）
    :return: 文件扩展名（包括点）
    """
    try:
        image = Image.open(io.BytesIO(image_data))
        return f".{image.format.lower()}"
    except Exception:
        return ".jpg"  # 默认扩展名

def process_content(content, file_prefix, image_folder):
    """
    处理Markdown内容，下载图片并替换图片链接。
    
    :param content: Markdown文件内容
    :param file_prefix: Markdown文件名（不包括扩展名），用于生成图片名称
    :param image_folder: 图片存储文件夹路径
    :return: 处理后的内容和图片计数
    """
    image_count = 0

    def replace_image(match):
        nonlocal image_count
        image_count += 1
        image_url = match.group(1)
        
        print(f"处理图片 #{image_count}: {image_url}")
        
        # 下载图片
        image_data = download_image(image_url)
        if image_data is None:
            print(f"无法下载图片，保留原始链接: {image_url}")
            return match.group()  # 如果下载失败，保持原始链接
        
        # 获取文件扩展名
        file_extension = get_file_extension(image_data)
        
        new_image_name = f"{file_prefix}_{image_count}{file_extension}"
        full_image_path = os.path.join(image_folder, new_image_name)
        
        # 保存图片
        try:
            with open(full_image_path, 'wb') as f:
                f.write(image_data)
            print(f"成功保存图片: {full_image_path}")
            return f'![](./images/{new_image_name})'
        except Exception as e:
            print(f"保存图片失败，保留原始链接: {str(e)}")
            return match.group()  # 如果保存失败，保持原始链接

    # 使用正则表达式查找并替换图片链接
    new_content = re.sub(r'!\[.*?\]\((http[s]?://[^\s\)]+)\)', replace_image, content)
    
    return new_content, image_count

def process_single_file(filename, input_folder, output_folder, image_folder):
    """
    处理单个Markdown文件。
    
    :param filename: Markdown文件名
    :param input_folder: 输入文件夹路径
    :param output_folder: 输出文件夹路径
    :param image_folder: 图片存储文件夹路径
    """
    input_path = os.path.join(input_folder, filename)
    output_path = os.path.join(output_folder, filename)
    
    with open(input_path, 'r', encoding='utf-8') as file:
        content = file.read()
    
    # 使用文件名（不包括扩展名）作为图片名称的前缀
    file_prefix = os.path.splitext(filename)[0]
    new_content, _ = process_content(content, file_prefix, image_folder)
    
    # 写入处理后的内容到新文件
    with open(output_path, 'w', encoding='utf-8') as file:
        file.write(new_content)

def process_markdown_files(input_folder, output_folder, image_folder):
    """
    处理指定文件夹中的所有Markdown文件。
    
    :param input_folder: 输入Markdown文件的文件夹路径
    :param output_folder: 处理后Markdown文件的输出文件夹路径
    :param image_folder: 下载图片的存储文件夹路径
    """
    # 确保输出文件夹和图片文件夹存在
    os.makedirs(output_folder, exist_ok=True)
    os.makedirs(image_folder, exist_ok=True)

    # 遍历输入文件夹中的所有Markdown文件
    for filename in os.listdir(input_folder):
        if filename.endswith('.md'):
            try:
                process_single_file(filename, input_folder, output_folder, image_folder)
                print(f"成功处理文件: {filename}")
            except Exception as e:
                print(f"处理文件 {filename} 时出错: {str(e)}")

if __name__ == "__main__":
    input_folder = 'D:\\学习\\markdown文档\\技术文章摘抄\\全解网络协议'
    output_folder = os.path.join(input_folder, 'local')
    image_folder = os.path.join(output_folder, 'images')

    process_markdown_files(input_folder, output_folder, image_folder)
